package com.example.config;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import java.util.List;

/**
 * @author Blue
 * @version 1.0
 */
@Configuration
public class ChatConfig {

    @Bean
    public ChatMemoryProvider chatMemoryProvider(){
        return new ChatMemoryProvider() {
            @Override
            public ChatMemory get(Object memoryId) {
                return MessageWindowChatMemory.builder().id(memoryId).maxMessages(10).build();
            }
        };
    }

    @Primary
    @Bean
    public EmbeddingStore embeddingStore(){
        // 加载文件
        Document doc = ClassPathDocumentLoader.loadDocument("README.md");
        // 构建向量数据库对象
        InMemoryEmbeddingStore store = new InMemoryEmbeddingStore();
        // 构建文本数据切割对象
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder().embeddingStore(store).build();
        ingestor.ingest(doc);
        return store;

    }

    @Bean
    public ContentRetriever contentRetriever(EmbeddingStore store){
        // 构建向量数据库检索对象
        return EmbeddingStoreContentRetriever.builder()
            .embeddingStore(store).maxResults(10).minScore(0.6).build();
    }
}
